"""
-*- coding: utf-8 -*-
@File  : generate_frames.py
@author: ZhenyuYang
@Time  : 2023/03/30 19:38
"""
import pandas as pd
import os
import cv2
import re

dataset = ["casme2", "smic", "samm"]
face_split = False
annotation_path = "../combined_3class_gt.csv"
annotations_frame = pd.read_csv(annotation_path, header=None)
annotations = annotations_frame.to_numpy()
print("Annotation shape: {}".format(annotations.shape))
labels = ["negative", "positive", "surprise"]
if face_split:
    import dlib
    detector = dlib.get_frontal_face_detector()
if not os.path.exists("../data/flow_prepare_rawframes"):
    os.mkdir("../data/flow_prepare_rawframes")

count = 0   # 已生成的样本数
for annotation in annotations:
    # print(annotation)
    count += 1
    if annotation[0] not in dataset:
        continue
    if annotation[0] == "casme2":
        data_folder = "../data/CASME2/{}/{}".format(annotation[1], annotation[2])
    elif annotation[0] == "smic":
        data_info = annotation[2].split("_")
        sub_num = int(data_info[0][1:])
        data_folder = "../data/SMIC/s{}/micro/{}/s{}_{}_{}".format(
            sub_num, labels[annotation[3]], sub_num, data_info[1], data_info[2])
    elif annotation[0] == "samm":
        data_folder = "../data/SAMM/{}/{}".format(annotation[1], annotation[2])
    else:
        continue
    data_label = labels[annotation[3]]
    # print(data_label)
    data_list = os.listdir(data_folder)
    # print(data_list)
    if not os.path.exists("../data/flow_prepare_rawframes/{}".format(data_label)):
        os.mkdir("../data/flow_prepare_rawframes/{}".format(data_label))
    dst_folder = "../data/flow_prepare_rawframes/{}/{}-{}-{}/".format(
        data_label, annotation[0], annotation[1], annotation[2])
    if not os.path.exists(dst_folder):
        os.mkdir(dst_folder)
    print("Generating \'{}\', {} frames, {}/{}".format(dst_folder, len(data_list), count, annotations.shape[0]))
    start_count = 999999999
    for img_name in data_list:
        img_count = int(re.findall(r"\d+", img_name)[0])
        if img_count < start_count:
            start_count = img_count
    for img_name in data_list:
        img = cv2.imread("{}/{}".format(data_folder, img_name))
        # cv2.imshow('Raw_image', img)
        # cv2.waitKey(0)
        if annotation[0] == "casme2":
            if face_split:
                face_position = detector(img)
                face = img[face_position[0].left(): face_position[0].right(),
                           face_position[0].top(): face_position[0].bottom(), :]
            else:
                face = img
            # cv2.imshow('Face', face)
            # cv2.waitKey(0)
        elif annotation[0] == "smic":
            face = img
        elif annotation[0] == "samm":
            face = img
        else:
            break
        img_count = int(re.findall(r"\d+", img_name)[0]) - start_count + 1
        if img_count <= 0:
            raise Exception()
        cv2.imwrite("{}{:05d}.jpg".format(dst_folder, img_count), face)
